from comman.activator import *
from comman.layers import *
from comman.net import Sequential
from comman.optimizers import *
from dataset.mnist.mnist import load_mnist
from util.image_util import *

epochs = 20
batch_size = 100
learning_rate = 1e-1

# 读入数据
(x_train, t_train), (x_test, t_test) = load_mnist(one_hot_label=True)
x_validation, t_validation = x_test[:batch_size], t_test[:batch_size]
# x_test, t_test = x_test[batch_size:], t_test[batch_size:]

net = Sequential([
    Dense(200, Relu()),
    Dropout(0.3),
    Dense(50, Relu()),
    Dropout(0.3),
    Dense(10, None),
])

net.compile(SoftmaxWithLoss(), SGD(learning_rate))

net.fit(x_train, t_train, epochs=epochs, batch_size=batch_size,
        validation_data=(x_validation, t_validation))

accuracy = net.evaluate(x_test, t_test)
print(f"accuracy={accuracy}")

zoom = 50
for i in range(x_test.shape[0]):
    x = x_test[i]
    t = t_test[i]
    y = net.predict(x)
    y = np.argmax(y)
    if t[y] != 1:
        # 把图像的形状变成原来的尺寸
        img_arr = x.reshape(28, 28)
        img = convert_img(img_arr * 255)
        img = img_zoom(img, zoom=zoom)
        pic_text(img, (10, 10), str(y), fill="#ffffff")
        pic_text(img, (28 * zoom - 100, 28 * zoom - 100), str(np.argmax(t)), fill="#ffffff")
        print(f"{y} --> {np.argmax(t)}")
        img.show()
